{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<center>\n",
    "    <p style=\"text-align:center\">\n",
    "        <img alt=\"phoenix logo\" src=\"https://storage.googleapis.com/arize-phoenix-assets/assets/phoenix-logo-light.svg\" width=\"200\"/>\n",
    "        <br>\n",
    "        <a href=\"https://arize.com/docs/phoenix/\">Docs</a>\n",
    "        |\n",
    "        <a href=\"https://github.com/Arize-ai/phoenix\">GitHub</a>\n",
    "        |\n",
    "        <a href=\"https://arize-ai.slack.com/join/shared_invite/zt-2w57bhem8-hq24MB6u7yE_ZF_ilOYSBw#/shared-invite/email\">Community</a>\n",
    "    </p>\n",
    "</center>\n",
    "\n",
    "# Google GenAI SDK - Building a Routing Agent"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q google-genai arize-phoenix-otel openinference-instrumentation-google-genai"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Connect to Arize Phoenix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from getpass import getpass\n",
    "\n",
    "from google import genai\n",
    "\n",
    "from phoenix.otel import register\n",
    "\n",
    "if \"PHOENIX_API_KEY\" not in os.environ:\n",
    "    os.environ[\"PHOENIX_API_KEY\"] = getpass(\"🔑 Enter your Phoenix API key: \")\n",
    "\n",
    "if \"PHOENIX_COLLECTOR_ENDPOINT\" not in os.environ:\n",
    "    os.environ[\"PHOENIX_COLLECTOR_ENDPOINT\"] = getpass(\"🔑 Enter your Phoenix Collector Endpoint\")\n",
    "\n",
    "tracer_provider = register(auto_instrument=False, project_name=\"google-genai-routing-agent\")\n",
    "tracer = tracer_provider.get_tracer(__name__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Authenticate with Google Vertex AI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!gcloud auth login"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from google.genai import types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a client using the Vertex AI API, you could also use the Google GenAI API instead here\n",
    "client = genai.Client(vertexai=True, project=\"<ADD YOUR GCP PROJECT ID>\", location=\"us-central1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Routing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tracer.tool()\n",
    "def get_current_weather(location: str) -> str:\n",
    "    \"\"\"Returns the current weather.\n",
    "\n",
    "    Args:\n",
    "      location: The city and state, e.g. San Francisco, CA\n",
    "    \"\"\"\n",
    "    return \"sunny\"\n",
    "\n",
    "\n",
    "@tracer.tool()\n",
    "def get_current_time(location: str) -> str:\n",
    "    \"\"\"Returns the current time.\n",
    "\n",
    "    Args:\n",
    "      location: The city and state, e.g. San Francisco, CA\n",
    "    \"\"\"\n",
    "    return \"10:00 AM\"\n",
    "\n",
    "\n",
    "@tracer.llm()\n",
    "def call_google_genai(query: str):\n",
    "    response = client.models.generate_content(\n",
    "        model=\"gemini-2.0-flash-001\",\n",
    "        contents=query,\n",
    "        config=types.GenerateContentConfig(tools=[get_current_weather, get_current_time]),\n",
    "    )\n",
    "    return response.text\n",
    "\n",
    "\n",
    "print(call_google_genai(query=\"What is the weather like in Boston?\"))\n",
    "print(call_google_genai(query=\"What is the time in Boston?\"))"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
